#!/usr/bin/env python
# -*- coding: utf-8 -*-
# author: tangzhichao

'''
语义切片
'''

from typing import List
from .base import BaseSlicer
# from langchain_community.embeddings import HuggingFaceEmbeddings


class SemanticSlicer(BaseSlicer):
    """基于语义的文档切片"""

    def __init__(self, chunk_size: int = 500, model_name: str = "BAAI/bge-small-zh"):
        self.embeddings = HuggingFaceEmbeddings(model_name=model_name)
        self.splitter = SemanticChunker(
            embeddings=self.embeddings,
            breakpoint_threshold_type="percentile",
            chunk_size=chunk_size
        )

    def split(self, text: str) -> List[str]:
        return self.splitter.split_text(text)

    @property
    def strategy_name(self) -> str:
        return "semantic"